from fastapi import FastAPI, File, UploadFile, HTTPException, Depends, BackgroundTasks
from fastapi.middleware.cors import CORSMiddleware
from fastapi.responses import JSONResponse
from pydantic import BaseModel, Field
import os
import uvicorn
import shutil
import uuid
from typing import List, Optional

# 导入自定义模块
from vector_store.vector_store import VectorStoreManager
from llm.llm_wrapper import LLMWrapper
from database.mysql_manager import db_manager
from core.order_knowledge_base import get_order_knowledge_base
from fastapi import FastAPI, File, UploadFile, HTTPException, BackgroundTasks

# 初始化FastAPI应用
bapp = FastAPI(title="RAG智能检索系统", version="1.0.0")

# 添加CORS中间件以允许跨域请求
bapp.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],  # 在生产环境中应限制为特定的源
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)

# 创建数据目录用于存储上传的文件和向量数据库
DATA_DIR = "data"
UPLOAD_DIR = os.path.join(DATA_DIR, "uploads")
if not os.path.exists(DATA_DIR):
    os.makedirs(DATA_DIR)
if not os.path.exists(UPLOAD_DIR):
    os.makedirs(UPLOAD_DIR)

# 初始化MySQL数据库
print("正在初始化MySQL数据库...")
if db_manager.initialize_database():
    print("MySQL数据库初始化成功")
    # 插入初始数据
    if db_manager.insert_initial_data():
        print("MySQL初始数据插入成功")
    else:
        print("MySQL初始数据插入失败，但不影响系统运行")
else:
    print("MySQL数据库初始化失败，请检查配置")

# 初始化订单知识库
try:
    order_knowledge_base = get_order_knowledge_base()
    print("订单知识库初始化成功")
except Exception as e:
    print(f"初始化订单知识库时出错: {e}")
    order_knowledge_base = None

# 初始化向量存储管理器
try:
    vector_store_manager = VectorStoreManager()
    print("向量存储组件已成功初始化")
except Exception as e:
    print(f"初始化向量存储组件时出错: {e}")
    vector_store_manager = None

# 初始化LLM包装器
try:
    llm_wrapper = LLMWrapper()
    print("LLM包装器初始化完成")
except Exception as e:
    print(f"初始化LLM包装器时出错: {e}")
    llm_wrapper = None

# 定期更新订单知识库的函数
def update_order_knowledge_task():
    """定期更新订单知识库的后台任务"""
    if order_knowledge_base and vector_store_manager:
        try:
            print("开始定期更新订单知识库...")
            success = order_knowledge_base.update_order_knowledge_base(vector_store_manager)
            if success:
                print("订单知识库更新成功")
            else:
                print("订单知识库更新失败")
        except Exception as e:
            print(f"更新订单知识库时出错: {e}")

# 启动后台任务定期更新订单知识库（在应用启动后运行一次）
import asyncio
import threading

async def start_periodic_update():
    # 延迟20秒执行第一次更新，给其他组件充分的初始化时间
    await asyncio.sleep(20)
    # 使用线程池执行阻塞操作
    loop = asyncio.get_event_loop()
    await loop.run_in_executor(None, update_order_knowledge_task)

# 在应用启动时启动定期更新订单知识库的任务
@bapp.on_event("startup")
async def startup_event():
    print("应用启动中，准备启动订单知识库定期更新任务...")
    # 创建一个任务来运行定期更新
    asyncio.create_task(start_periodic_update())

# 基本的健康检查端点
@bapp.get("/")
def read_root():
    return {"status": "success", "message": "RAG智能检索系统已启动", "version": "1.0.0"}

# 定义请求模型
class QueryModel(BaseModel):
    query: str
    collection_name: str
    top_k: int = 3
    use_llm: bool = True

class CollectionModel(BaseModel):
    collection_name: str = Field(..., min_length=1, max_length=100, pattern=r'^[a-zA-Z0-9_]+$')
    description: Optional[str] = None

# 获取所有集合
@bapp.get("/collections")
def get_collections():
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    try:
        collections = vector_store_manager.get_collections()
        print(f"即将返回的集合列表: {collections}")
        response = {"status": "success", "collections": collections}
        print(f"返回的响应: {response}")
        return response
    except Exception as e:
        print(f"获取集合时发生异常: {e}")
        raise HTTPException(status_code=500, detail=f"获取集合时出错: {str(e)}")

# 创建集合
@bapp.post("/collections")
def create_collection(collection: CollectionModel):
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    try:
        # 检查集合名称是否合法
        if not collection.collection_name.isalnum() and '_' not in collection.collection_name:
            raise HTTPException(status_code=400, detail="集合名称只能包含字母、数字和下划线")
        
        # 检查集合是否已存在
        existing_collections = vector_store_manager.get_collections()
        if collection.collection_name in existing_collections:
            raise HTTPException(status_code=400, detail=f"集合 '{collection.collection_name}' 已存在")
        
        # 实际创建集合
        try:
            # 使用一个空文档来触发集合的实际创建
            # 这是因为ChromaDB需要在集合中至少有一个文档
            # 这里我们使用一个临时的空文档来创建集合
            from langchain.schema import Document
            empty_doc = [Document(page_content="", metadata={"source": "collection_creation_marker"})]
            # 调用add_documents来创建集合
            vector_store_manager.add_documents(collection.collection_name, empty_doc)
            print(f"成功创建集合 '{collection.collection_name}'")
        except Exception as create_error:
            print(f"创建集合时出错: {create_error}")
            raise HTTPException(status_code=500, detail=f"创建集合时出错: {str(create_error)}")
        
        return {"status": "success", "message": f"集合 '{collection.collection_name}' 创建成功"}
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"创建集合时出错: {str(e)}")

# 删除集合
@bapp.delete("/collections/{collection_name}")
def delete_collection(collection_name: str):
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    try:
        result = vector_store_manager.delete_collection(collection_name)
        if result:
            return {"status": "success", "message": f"集合 '{collection_name}' 删除成功"}
        else:
            raise HTTPException(status_code=404, detail=f"集合 '{collection_name}' 不存在")
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"删除集合时出错: {str(e)}")

# 获取集合中的所有文档
@bapp.get("/collections/{collection_name}/documents")
def get_collection_documents(collection_name: str):
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    try:
        documents = vector_store_manager.get_all_documents_in_collection(collection_name)
        return {"status": "success", "documents": documents}
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"获取集合中文档时出错: {str(e)}")

# 删除集合中的单个文档
@bapp.delete("/collections/{collection_name}/documents/{document_id}")
def delete_document(collection_name: str, document_id: str):
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    try:
        result = vector_store_manager.delete_document_from_collection(collection_name, document_id)
        if result:
            return {"status": "success", "message": f"文档 '{document_id}' 删除成功"}
        else:
            raise HTTPException(status_code=404, detail=f"文档 '{document_id}' 不存在或删除失败")
    except HTTPException as e:
        raise e
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"删除文档时出错: {str(e)}")

# 上传文档到集合
@bapp.post("/collections/{collection_name}/documents")
async def upload_documents(
    collection_name: str,
    files: List[UploadFile] = File(...)
):
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    
    try:
        uploaded_files = []
        total_chunks = 0
        
        for file in files:
            # 生成唯一的文件名
            file_extension = os.path.splitext(file.filename)[1]
            unique_filename = f"{uuid.uuid4()}{file_extension}"
            file_path = os.path.join(UPLOAD_DIR, unique_filename)
            
            # 保存文件
            with open(file_path, "wb") as buffer:
                shutil.copyfileobj(file.file, buffer)
            
            # 加载文档
            documents = vector_store_manager.load_document(file_path)
            
            # 添加到向量存储
            chunks_count = vector_store_manager.add_documents(collection_name, documents)
            total_chunks += chunks_count
            
            uploaded_files.append({
                "filename": file.filename,
                "saved_as": unique_filename,
                "chunks": chunks_count
            })
            
        return {
            "status": "success",
            "message": f"成功上传 {len(files)} 个文件到集合 '{collection_name}'",
            "uploaded_files": uploaded_files,
            "total_chunks": total_chunks
        }
    
    except ValueError as e:
        raise HTTPException(status_code=400, detail=str(e))
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"上传文档时出错: {str(e)}")

# 智能搜索端点
@bapp.post("/search")
def smart_search(query: QueryModel):
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    
    try:
        # 在向量存储中搜索相关文档
        search_results = vector_store_manager.search(query.collection_name, query.query, query.top_k)
        
        if not search_results:
            return {
                "status": "success",
                "query": query.query,
                "message": "未找到相关文档",
                "results": [],
                "answer": "基于当前知识库，我无法回答这个问题。"
            }
        
        # 如果需要使用LLM生成回答
        if query.use_llm and llm_wrapper is not None:
            answer_result = llm_wrapper.generate_answer(query.query, search_results)
        else:
            answer_result = {"answer": "请查看相关文档片段以获取答案。"}
        
        return {
            "status": "success",
            "query": query.query,
            "results": search_results,
            "answer": answer_result["answer"],
            "model_info": answer_result.get("model_used") if query.use_llm else None,
            "token_usage": answer_result.get("token_usage") if query.use_llm else None
        }
    
    except Exception as e:
        raise HTTPException(status_code=500, detail=f"搜索时出错: {str(e)}")

# 订单知识库的请求模型
class OrderKnowledgeQueryModel(BaseModel):
    query: str
    top_k: int = 3
    use_llm: bool = True

# 订单相关的请求模型
class OrderItemModel(BaseModel):
    product_name: str = Field(..., min_length=1, max_length=100)
    quantity: int = Field(..., gt=0)
    unit_price: float = Field(..., ge=0.01)

class CreateOrderModel(BaseModel):
    order_no: str = Field(..., min_length=1, max_length=50, pattern=r'^[a-zA-Z0-9-_]+$')
    customer_name: str = Field(..., min_length=1, max_length=100)
    customer_email: str = Field(..., min_length=1, max_length=100)
    items: List[OrderItemModel]
    status: str = Field(default="pending", pattern=r'^(pending|completed|canceled)$')

class UpdateOrderStatusModel(BaseModel):
    status: str = Field(..., pattern=r'^(pending|completed|canceled)$')

# 获取所有订单
@bapp.get("/orders")
def get_orders():
    try:
        # 查询所有订单
        orders = db_manager.execute_query("""
            SELECT id, order_no, customer_name, customer_email, total_amount, status,
                   created_at, updated_at
            FROM orders
            ORDER BY created_at DESC
        """)
        
        return {
            "status": "success",
            "data": orders,
            "count": len(orders)
        }
    except Exception as e:
        print(f"获取订单列表时出错: {e}")
        raise HTTPException(status_code=500, detail=f"获取订单列表时出错: {str(e)}")

# 获取订单详情
@bapp.get("/orders/{order_id}")
def get_order_detail(order_id: int):
    try:
        # 查询订单基本信息
        order = db_manager.execute_query("""
            SELECT id, order_no, customer_name, customer_email, total_amount, status,
                   created_at, updated_at
            FROM orders
            WHERE id = %s
        """, (order_id,))
        
        if not order:
            raise HTTPException(status_code=404, detail=f"订单ID {order_id} 不存在")
        
        # 查询订单商品列表
        items = db_manager.execute_query("""
            SELECT id, product_name, quantity, unit_price
            FROM order_items
            WHERE order_id = %s
        """, (order_id,))
        
        order[0]["items"] = items
        
        return {
            "status": "success",
            "data": order[0]
        }
    except HTTPException as e:
        raise e
    except Exception as e:
        print(f"获取订单详情时出错: {e}")
        raise HTTPException(status_code=500, detail=f"获取订单详情时出错: {str(e)}")

# 创建订单
@bapp.post("/orders")
def create_order(order: CreateOrderModel):
    try:
        # 检查订单号是否已存在
        existing_order = db_manager.execute_query(
            "SELECT id FROM orders WHERE order_no = %s", (order.order_no,)
        )
        if existing_order:
            raise HTTPException(status_code=400, detail=f"订单号 {order.order_no} 已存在")
        
        # 计算订单总金额
        total_amount = sum(item.quantity * item.unit_price for item in order.items)
        
        # 插入订单基本信息
        order_id = db_manager.execute_query(
            """
            INSERT INTO orders (order_no, customer_name, customer_email, total_amount, status)
            VALUES (%s, %s, %s, %s, %s)
            RETURNING id
            """,
            (order.order_no, order.customer_name, order.customer_email, total_amount, order.status)
        )
        
        if not order_id:
            raise Exception("创建订单失败，无法获取订单ID")
        
        order_id = order_id[0]["id"]
        
        # 插入订单商品
        order_items = []
        for item in order.items:
            order_items.append((
                order_id,
                item.product_name,
                item.quantity,
                item.unit_price
            ))
        
        db_manager.execute_many(
            """
            INSERT INTO order_items (order_id, product_name, quantity, unit_price)
            VALUES (%s, %s, %s, %s)
            """,
            order_items
        )
        
        # 获取创建的订单详情
        created_order = db_manager.execute_query(
            """
            SELECT id, order_no, customer_name, customer_email, total_amount, status,
                   created_at, updated_at
            FROM orders
            WHERE id = %s
            """,
            (order_id,)
        )
        
        return {
            "status": "success",
            "message": "订单创建成功",
            "data": created_order[0]
        }
    except HTTPException as e:
        raise e
    except Exception as e:
        print(f"创建订单时出错: {e}")
        raise HTTPException(status_code=500, detail=f"创建订单时出错: {str(e)}")

# 更新订单状态
@bapp.put("/orders/{order_id}/status")
def update_order_status(order_id: int, status_data: UpdateOrderStatusModel):
    try:
        # 检查订单是否存在
        existing_order = db_manager.execute_query(
            "SELECT id FROM orders WHERE id = %s", (order_id,)
        )
        if not existing_order:
            raise HTTPException(status_code=404, detail=f"订单ID {order_id} 不存在")
        
        # 更新订单状态
        affected_rows = db_manager.execute_update(
            "UPDATE orders SET status = %s WHERE id = %s",
            (status_data.status, order_id)
        )
        
        if affected_rows > 0:
            # 获取更新后的订单信息
            updated_order = db_manager.execute_query(
                """
                SELECT id, order_no, customer_name, customer_email, total_amount, status,
                       created_at, updated_at
                FROM orders
                WHERE id = %s
                """,
                (order_id,)
            )
            
            return {
                "status": "success",
                "message": "订单状态更新成功",
                "data": updated_order[0]
            }
        else:
            raise Exception("订单状态更新失败")
    except HTTPException as e:
        raise e
    except Exception as e:
        print(f"更新订单状态时出错: {e}")
        raise HTTPException(status_code=500, detail=f"更新订单状态时出错: {str(e)}")

# 删除订单
@bapp.delete("/orders/{order_id}")
def delete_order(order_id: int):
    try:
        # 检查订单是否存在
        existing_order = db_manager.execute_query(
            "SELECT id FROM orders WHERE id = %s", (order_id,)
        )
        if not existing_order:
            raise HTTPException(status_code=404, detail=f"订单ID {order_id} 不存在")
        
        # 删除订单（会级联删除订单项）
        affected_rows = db_manager.execute_update(
            "DELETE FROM orders WHERE id = %s",
            (order_id,)
        )
        
        if affected_rows > 0:
            return {
                "status": "success",
                "message": "订单删除成功"
            }
        else:
            raise Exception("订单删除失败")
    except HTTPException as e:
        raise e
    except Exception as e:
        print(f"删除订单时出错: {e}")
        raise HTTPException(status_code=500, detail=f"删除订单时出错: {str(e)}")

# 查询订单知识库
@bapp.post("/order-knowledge/search")
def search_order_knowledge(query: OrderKnowledgeQueryModel):
    """
    在订单知识库中搜索相关信息
    支持自然语言查询，并可选择是否使用LLM生成回答
    """
    if order_knowledge_base is None:
        raise HTTPException(status_code=500, detail="订单知识库未初始化")
    
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    
    try:
        # 在订单知识库中搜索相关文档
        search_results = order_knowledge_base.search_order_knowledge(
            vector_store_manager, 
            query.query, 
            query.top_k
        )
        
        if not search_results:
            return {
                "status": "success",
                "query": query.query,
                "message": "未找到相关订单信息",
                "results": [],
                "answer": "基于当前订单知识库，我无法回答这个问题。"
            }
        
        # 如果需要使用LLM生成回答
        if query.use_llm and llm_wrapper is not None:
            answer_result = order_knowledge_base.generate_order_answer(
                llm_wrapper, 
                query.query, 
                search_results
            )
        else:
            answer_result = {"answer": "请查看相关订单信息以获取答案。"}
        
        return {
            "status": "success",
            "query": query.query,
            "results": search_results,
            "answer": answer_result["answer"],
            "model_info": answer_result.get("model_used") if query.use_llm else None,
            "token_usage": answer_result.get("token_usage") if query.use_llm else None
        }
        
    except Exception as e:
        print(f"搜索订单知识库时出错: {e}")
        raise HTTPException(status_code=500, detail=f"搜索订单知识库时出错: {str(e)}")

# 手动更新订单知识库
@bapp.post("/order-knowledge/update")
def manual_update_order_knowledge(background_tasks: BackgroundTasks):
    """
    手动触发订单知识库更新
    此操作会从MySQL重新加载订单数据并更新向量存储
    """
    if order_knowledge_base is None:
        raise HTTPException(status_code=500, detail="订单知识库未初始化")
    
    if vector_store_manager is None:
        raise HTTPException(status_code=500, detail="向量存储组件未初始化")
    
    # 使用后台任务执行更新，避免请求超时
    background_tasks.add_task(update_order_knowledge_task)
    
    return {
        "status": "success",
        "message": "订单知识库更新任务已启动，请稍后查看更新结果"
    }

if __name__ == "__main__":
    uvicorn.run(bapp, host="0.0.0.0", port=8000)